import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
import os

# create dir
def pidToid(csv):
    pid_to_id = {}
    counter = 0
    for pid in csv['Pid']:
        if pid not in pid_to_id:
            pid_to_id[pid] = counter
            counter += 1
    return pid_to_id

def crateGraph(csv):
    dirPid = pidToid(csv)
    G= nx.DiGraph()
    graph_call = []
    log_time = 0
    for index,row in csv.iterrows():
        graphId = dirPid[row["Pid"]]
        nowGraphcall = row['Syscall'] + str(graphId)
        #  新新图
        if len(graph_call) < graphId + 1 :
            if len(graph_call) == 0:
                G.add_node(nowGraphcall, time = log_time , graphNum = graphId)
                log_time +=1
                graph_call.append(nowGraphcall)
                # 旧的新图
            else:
                G.add_node(nowGraphcall, time = log_time , graphNum = graphId)
                G.add_edge(oldGraphcall,nowGraphcall)
                log_time +=1
                graph_call.append(nowGraphcall)
                
        else:
            # 旧的同一图
            if graphId == oldGraphId :
                G.add_node(nowGraphcall, time = log_time , graphNum = graphId)
                G.add_edge(oldGraphcall,nowGraphcall)
                log_time +=1
                graph_call[graphId] = nowGraphcall
            # 旧的非同图
            else:
                G.add_node(nowGraphcall, time = log_time , graphNum = graphId)
                G.add_edge(oldGraphcall,nowGraphcall)
                G.add_edge(graph_call[graphId],nowGraphcall)
                log_time +=1
                graph_call[graphId] = nowGraphcall
        oldGraphId = graphId
        oldGraphcall = graph_call[graphId]
    num_nodes = G.number_of_nodes()
    num_edges = G.number_of_edges()
    print(f"The graph has {num_nodes} nodes and {num_edges} edges.")
    return G

# # 使用spring布局
# pos = nx.spring_layout(G)
# # 绘制图形
# plt.figure(figsize=(8, 6))  # 设置图形大小
# nx.draw(G, pos, with_labels=True, node_size=1500, node_color="skyblue", font_size=15, width=2, edge_color='gray')
# plt.show()
# plt.savefig('test.png')

# nx.write_gpickle(G, "/root/jx_codes/crateGraph/graph/test.gpickle")

directory = '/root/jx_codes/crateGraph/csv'
csv_files = [f for f in os.listdir(directory) if f.endswith('.csv')]
for csv_file in csv_files:
    file_path = os.path.join(directory, csv_file)
    csv = pd.read_csv(file_path)
    G = crateGraph(csv)
    exec(f'nx.write_gpickle(G, "/root/jx_codes/crateGraph/graph/{csv_file.replace(".csv","")}.gpickle")')